##########################################################################################

library(data.table)
library(optparse)
library(Seurat)
library(ArchR)
library(ggsci)
library(dplyr)

##########################################################################################
option_list <- list(
    make_option(c("--comine_data_file"), type = "character"),
    make_option(c("--nclust"), type = "character"),
    make_option(c("--scriptPath"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 整合atac和rna的文件
    comine_data_file <- "~/20231121_singleMuti/results/qc_atac_v2/germ/testis_combined_peak.combineRNA.qc.Rdata"

    ## 既往研究整理的代码
    scriptPath <- "~/20231121_singleMuti/scripts/scScalpChromatin"

    ## 
    nclust <- 11

    ## 输出
    out_path <- "~/20231121_singleMuti/results/celltype_plot/peak2gene"

}


###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

comine_data_file <- opt$comine_data_file
scriptPath <- opt$scriptPath
nclust <- as.numeric(opt$nclust)
out_path <- opt$out_path

dir.create( out_path , recursive = T)

##########################################################################################
## 已发表文献写好的脚本
source(paste0(scriptPath, "/plotting_config.R"))
source(paste0(scriptPath, "/misc_helpers.R"))
source(paste0(scriptPath, "/matrix_helpers.R"))
source(paste0(scriptPath, "/archr_helpers.R"))
source(paste0(scriptPath, "/GO_wrappers.R"))

##########################################################################################
## 导入数据
a <- load(comine_data_file)
# testis_combined_peak_combineRNA

## 细胞顺序
cell_order <- c("SSC", "Differenting&Differented SPG", "Leptotene",
    "Zygotene", "Patchytene", "Diplotene",
    "Early stage of spermatids", "Round&ElongateS.tids", "Sperm",
    "Leydig cells", "Myoid cells", "Pericytes",
    "Sertoli cells", "Endothelial cells", "NKT cells", "Macrophages"
    )

## 细胞颜色
use_colors <- c(pal_npg("nrc")(10) , pal_jco("default")(6))
names(use_colors) <- c("Myoid cells" , "Leydig cells" , "Endothelial cells" , "Zygotene" , "Round&ElongateS.tids" , 
"Patchytene" , "SSC" , "Sperm" , "Diplotene" , "Early stage of spermatids" , "Leptotene" , 
"Sertoli cells" , "Macrophages" , "Differenting&Differented SPG" , "Pericytes" , "NKT cells" )

##########################################################################################
## To identify peak-to-gene links in ArchR
projHeme5 <- testis_combined_peak_combineRNA
projHeme5 <- addPeak2GeneLinks(ArchRProj = projHeme5 , useMatrix = "GeneExpressionMatrix")

## 参数参考
## https://github.com/GreenleafLab/scScalpChromatin/blob/9e333bd3194ed6548cfad35c4b9ba678b0cdde31/Figure_2_Linked_Peaks.R#L71
# P2G definition cutoffs
corrCutoff <- 0.5       # Default in plotPeak2GeneHeatmap is 0.45
varCutoffATAC <- 0.25   # Default in plotPeak2GeneHeatmap is 0.25
varCutoffRNA <- 0.25    # Default in plotPeak2GeneHeatmap is 0.25

########################################
## Plotting a heatmap of peak-to-gene links
#projHeme5@cellColData$cell_type <- factor( projHeme5@cellColData$cell_type , levels = cell_order , order = T )
#nclust <- 25 
p <- plotPeak2GeneHeatmap(
  projHeme5, 
  corCutOff = corrCutoff, 
  groupBy="cell_type", 
  nPlot = 1000000, returnMatrices=FALSE, 
  k=nclust, seed=1, palGroup=use_colors
  )

out_file <- paste0(out_path, sprintf("/peakToGeneHeatmap_LabelClust_k%s.pdf", nclust))
pdf(out_file , width=16, height=15)
draw(p)
dev.off()

########################################
## 计算每个cluster里面peak富集在哪些motif
atac_proj <- projHeme5

# Get all peaks
allPeaksGR <- getPeakSet(atac_proj)
allPeaksGR$peakName <- (allPeaksGR %>% {paste0(seqnames(.), "_", start(.), "_", end(.))})
names(allPeaksGR) <- allPeaksGR$peakName

# Need to force it to plot all peaks if you want to match the labeling when you 'returnMatrices'.
p2gMat <- plotPeak2GeneHeatmap(
  atac_proj, 
  corCutOff = corrCutoff, 
  groupBy="cell_type",
  nPlot = 1000000, returnMatrices=TRUE, 
  k=nclust, seed=1)

## 输出peak-gene相关参数
kclust_df_out <- cbind( kclust=p2gMat$ATAC$kmeansId , p2gMat$Peak2GeneLinks )
tmp_peak_dat <- data.frame(atac_proj@peakSet)
tmp_peak_dat$peak <- paste0( tmp_peak_dat$seqnames , ":" , tmp_peak_dat$start , "-" , tmp_peak_dat$end )
kclust_df_out <- merge( kclust_df_out , tmp_peak_dat , by = "peak" )

out_file <- paste0(out_path, sprintf("/peakToGeneHeatmap_LabelClust_k%s.tsv", nclust))
write.table(kclust_df_out , out_file , row.names = F , sep = "\t")

## 记录每个基因对应多少peak在各自的clust里面
peak_num <- data.frame(kclust_df_out) %>% 
group_by( gene , kclust ) %>%
summarize( peak_num = length(peak) )
out_file <- paste0(out_path, sprintf("/peakToGeneHeatmap_LabelClust_k%s.peakNum.tsv", nclust))
write.table(peak_num , out_file , row.names = F , sep = "\t")

# Get association of peaks to clusters
kclust_df <- data.frame(
  kclust=p2gMat$ATAC$kmeansId,
  peakName=p2gMat$Peak2GeneLinks$peak,
  gene=p2gMat$Peak2GeneLinks$gene
  )

# Fix peakname
kclust_df$peakName <- sapply(kclust_df$peakName, function(x) strsplit(x, ":|-")[[1]] %>% paste(.,collapse="_"))

# Get motif matches
matches <- getMatches(atac_proj, "Motif")
r1 <- SummarizedExperiment::rowRanges(matches)
rownames(matches) <- paste(seqnames(r1),start(r1),end(r1),sep="_")
matches <- matches[names(allPeaksGR)]

clusters <- unique(kclust_df$kclust) %>% sort()

enrichList <- lapply(clusters, function(x){
  cPeaks <- kclust_df[kclust_df$kclust == x,]$peakName %>% unique()
  ArchR:::.computeEnrichment(matches, which(names(allPeaksGR) %in% cPeaks), seq_len(nrow(matches)))
  }) %>% SimpleList
names(enrichList) <- clusters

# Format output to match ArchR's enrichment output
assays <- lapply(seq_len(ncol(enrichList[[1]])), function(x){
    d <- lapply(seq_along(enrichList), function(y){
        enrichList[[y]][colnames(matches),x,drop=FALSE]
      }) %>% Reduce("cbind",.)
    colnames(d) <- names(enrichList)
    d
  }) %>% SimpleList
names(assays) <- colnames(enrichList[[1]])
assays <- rev(assays)
res <- SummarizedExperiment::SummarizedExperiment(assays=assays)

formatEnrichMat <- function(mat, topN, minSig, clustCols=TRUE){
  plotFactors <- lapply(colnames(mat), function(x){
    ord <- mat[order(mat[,x], decreasing=TRUE),]
    ord <- ord[ord[,x]>minSig,]
    rownames(head(ord, n=topN))
  }) %>% do.call(c,.) %>% unique()
  pMat <- mat[plotFactors,]
  prettyOrderMat(pMat, clusterCols=clustCols)$mat
}

pMat <- formatEnrichMat(assays(res)$mlog10Padj, 5, 10, clustCols=FALSE)
# Save maximum enrichment
tfs <- strsplit(rownames(pMat), "_") %>% sapply(., `[`, 1)
rownames(pMat) <- paste0(tfs, " (", apply(pMat, 1, function(x) floor(max(x))), ")")

pMat <- apply(pMat, 1, function(x) x/max(x)) %>% t()
rownames(pMat) <- sapply( strsplit(rownames(pMat) , " ") , "[" , 1)

out_file <- paste0(out_path, sprintf("/enrichedMotifs_kclust_p2gHM_k%s.pdf", nclust))
pdf(out_file ,width=12, height=15)
ht_opt$simple_anno_size <- unit(0.25, "cm")
hm <- BORHeatmap(
  pMat, 
  limits=c(0,1), 
  clusterCols=FALSE, clusterRows=FALSE,
  labelCols=TRUE, labelRows=TRUE,
  dataColors = cmaps_BOR$comet,
  #top_annotation = ta,
  row_names_side = "left",
  width = ncol(pMat)*unit(0.5, "cm"),
  height = nrow(pMat)*unit(0.4, "cm"),
  border_gp=gpar(col="black"), # Add a black border to entire heatmap
  legendTitle="Norm.Enrichment -log10(P-adj)[0-Max]"
  )
draw(hm)
dev.off()

pMat <- formatEnrichMat(assays(res)$mlog10Padj, 5000, 10, clustCols=FALSE)
# Save maximum enrichment
tfs <- strsplit(rownames(pMat), "_") %>% sapply(., `[`, 1)
rownames(pMat) <- paste0(tfs, " (", apply(pMat, 1, function(x) floor(max(x))), ")")

pMat <- apply(pMat, 1, function(x) x/max(x)) %>% t()
rownames(pMat) <- sapply( strsplit(rownames(pMat) , " ") , "[" , 1)

out_file <- paste0(out_path, sprintf("/enrichedMotifs_kclust_p2gHM_k%s.tsv", nclust))
write.table(pMat , out_file , row.names = T , sep = "\t")

########################################
# GO enrichments of top N genes per cluster 
# ("Top" genes are defined as having the most peak-to-gene links)
kclust <- unique(kclust_df$kclust) %>% sort()
all_genes <- kclust_df$gene %>% unique() %>% sort()

# Save table of top linked genes per kclust
nGOgenes <- 200
topKclustGenes <- lapply(kclust, function(k){
  kclust_df[kclust_df$kclust == k,]$gene %>% getFreqs() %>% head(nGOgenes) %>% names()
  }) %>% do.call(cbind,.)
outfile <- paste0(out_path, sprintf("/topN_genes_kclust_k%s.tsv", nclust))
write.table(topKclustGenes, file=outfile, quote=FALSE, sep='\t', row.names = FALSE, col.names=TRUE)

# 通路富集
GOresults <- lapply(kclust, function(k){
  message(sprintf("Running GO enrichments on k cluster %s...", k))
  clust_genes <- topKclustGenes[,k]
  upGO <- rbind(
    calcTopGo(all_genes, interestingGenes=clust_genes, nodeSize=5, ontology="BP") 
    )

  upGO <- upGO[order(as.numeric(upGO$pvalue), decreasing=FALSE),]

  ## 构造GO对象
  geneList <- factor(as.integer(all_genes %in% clust_genes))
  names(geneList) <- all_genes
  GOdata <- suppressMessages(new(
    "topGOdata",
    ontology = "BP",
    allGenes = geneList,
    annot = annFUN.org, mapping = "org.Hs.eg.db", ID = "symbol",
    nodeSize = 5
  ))
  ## 提取每个通路中感兴趣的基因
  gene_in_pathway <- sapply(upGO$GO.ID, function(x)
    {
      genes <- genesInTerm(GOdata, x)
      # myGenes is the queried gene list
      paste0(genes[[1]][genes[[1]] %in% clust_genes] , collapse = "," )
    })
  
  upGO$gene_in_pathway <- gene_in_pathway
  upGO$cluster <- paste0("cluster" , k)

  return(upGO)
})

names(GOresults) <- paste0("cluster_", kclust)

## 输出为表格
pathway_res <- c()
for( clus in names(GOresults) ){

  tmp <- GOresults[[clus]]
  tmp$cluster <- clus
  pathway_res <- rbind( pathway_res , tmp )
}

out_file <- paste0(out_path, sprintf("/kclust_GO_3termsBPonlyBarLim_k%s.tsv", nclust))
write.table( pathway_res , out_file , row.names = F , quote = F , sep ="\t" )


# Plots of GO term enrichments:
pdf(paste0(out_path, sprintf("/kclust_GO_3termsBPonlyBarLim_k%s.pdf", nclust)), width=10, height=6)
for(name in names(GOresults)){
    goRes <- GOresults[[name]]
    if(nrow(goRes)>1){
      print(topGObarPlot(goRes, cmap = cmaps_BOR$comet, 
        nterms=10, border_color="black", 
        barwidth=0.85, title=name, barLimits=c(0, 15)))
    }
}
dev.off()

